'''
Created on May 4, 2010

@author: oabalbin
'''

import os
import sys
import glob
import shutil
import numpy as np
import pickle
from optparse import OptionParser
from datetime import datetime
from collections import deque, defaultdict


# My libraries
import signatures.parsers.read_microarray_data as pmd
import signatures.parsers.read_gene_lists as gpm

import signatures.db.gene_annotations as ga
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file


class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        self.tableDescription = 'probeID VARCHAR(40), geneSym VARCHAR(40), sample VARCHAR(40), log_fold FLOAT(20,10), pvalue FLOAT(20,10), INDEX probe (probeID), INDEX gene (geneSym), INDEX sample (sample), INDEX pvalue (pvalue)'
        
        # table with golub prostate cohort
        #self.tableName = 'ERG_ERG'
        self.tableName = 'mouseHuman'
        ######## Parameters ############
        # Parameters for pre-processing data:
        # percentage of low variable gene to be filter out of the analysis
        self.lvar=5 #15
        # percentage of low median gene to be filter out of the analysis
        self.lmed=5 #15
            
    
    def parse_input_annotfiles(self, annotInputFile):
    
        expr_file, seed_file, outdirs, db_file, samples, bfrm, param, typefile={},{},{},{},{},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '' or fields[0][0] == '#':
                continue
            
            if fields[0][0] == '@':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'@Array expression':expr_file, '@Seed genes':seed_file,'@Database file':db_file,'@Samples':samples,'@Output Directories':outdirs,\
                    '@BFRM':bfrm,'@Parameters':param}[typefile]
            if typefile=='@Parameters':
                dict[fields[0]] = fields[1].split(',')
            else:
                dict[fields[0]] = fields[1]
            
            
        return expr_file, seed_file, db_file, samples, bfrm, outdirs, param
    

    def read_files_folder(self,folderpath,ext):
        ''' '''
        # Read files in folder
        myfiles=[]
        for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
            myfiles.append(infile)
            #print "current file is: " + infile
        return myfiles
    
    
    def process_microarray_files(self,exp_files,tp,soft=False,matfile='', log10=False, mouse=False):
        dbf=[]
        for f in exp_files:
            f1 = f.split('/')[-1]
            of=f+'_dbfile'
            if f1 == 'nci60_RNA_Affy_HGU95_AE_RMA_sp.txt':
            #if f == '/home/oabalbin/projects/networks/exp_data_house/prostate_cell_lines_ERG/raw_data/nci60_RNA_Affy_HGU95_AE_RMA_sp.txt':
                print 'reading nci60 samples ' + f
                tp.microarray_nci60_file(open(f),open(of,'w'))
            elif soft:
                print 'reading soft samples ' + f
                sp=f+'_sampleslist'
                gene=f+'_genelist'
                spstatus = f+'_sampleslist_status'
                sp_tr=f+'_sampleslist_training'
                sp_tt=f+'_sampleslist_test'
                trstatus = f+'_sampleslist_training_status'
                ttstatus = f+'_sampleslist_test_status'
                
                tp.microarray_exp_softfile(open(f),open(of,'w'),open(gene,'w'), open(sp,'w'), open(spstatus,'w'), open(sp_tr,'w'), \
                                           open(sp_tt,'w'), open(trstatus,'w'), open(ttstatus,'w') )
            elif matfile:
                print 'reading a GEO matrix file ' + f
                tp.microarray_exp_geoMatrixfile(open(f), open(matfile), open(of,'w'))
            elif log10:
                print 'reading log10 samples ' + f
                tp.microarray_logexp_file(open(f),open(of,'w'),mouse)
                
            else:
                print 'reading house samples ' + f
                tp.microarray_exp_file(open(f),open(of,'w'))
            dbf.append(of)
        
        return dbf
        
    def create_database(self,dbfiles,doit=False):
        """
        It creates a databases if doit = True
        """
        if doit:
            print 'Making tables'            
            for dbf in dbfiles:     
                if not(self.mt.table_existence(self.tableName)):
                    self.mt.create_table(self.tableName,self.tableDescription)
                    
                self.mt.load_data_table(self.tableName,dbf)

                
    def initialize_database(self, exp_files,tp, doit=False, soft=False, matfile='', log10=False,mouse=False):
        """
        Initialize and create the tables in the database
        """
        dbfiles = self.process_microarray_files(exp_files,tp, soft, matfile, log10, mouse)
        self.create_database(dbfiles,doit)
    
    def check_create_dir( self, full_path_name ):
        """
        creates an output folder. Name folder is the date it was created
        returns the full path to the folder and the folder name only: outfolder, t.strftime
        """
        t = datetime.now()
        outfolder = full_path_name + t.strftime("%Y_%m_%d_%H_%M")
        if not os.path.isdir(outfolder):
            os.mkdir(outfolder)
        
        return outfolder, t.strftime("%Y_%m_%d_%H_%M")
         
         
    def write_arraytofile(self,outfile, H):
        """
        It writes the and numpy array to a file. 
        """
        for i in range(H.shape[0]):
            el = map(str,list(H[i,:]))
            outfile.write(",".join(el).replace(',','\t')+'\n')
            

def write_array2file(outfile, H, rownames=[]):
    """
    It writes the and numpy array to a file. 
    """
    for i in range(H.shape[0]):
        
        if np.ma.count_masked(H[i,:]) >= len(H[i,:]):
            continue
        
        el = map(str,list(H[i,:]))
        if rownames:
            outfile.write(str(rownames[i])+"\t"+",".join(el).replace(',','\t')+'\n')
        else:
            outfile.write(",".join(el).replace(',','\t')+'\n')




        

##########
if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", action="store_true", dest="dbtrue",
                            help="create a database. Bool")
    optionparser.add_option("-t", action="store_true", dest="softtrue",
                            help="use a microarray soft file. Bool")
    optionparser.add_option("-p", "--platformFile", dest="platformFile",
                            help="use a microarray soft file. Bool")
    # Running parameters
    optionparser.add_option("-u", action="store_true", dest="allgenes",
                            help="use genes in all uscs. Default use genes in universe file. Default universe. Bool")
    optionparser.add_option("-n", action="store_true", dest="norm",
                            help="normalize log2 and scale expression vals. Default only log2.  Bool")
    optionparser.add_option("-m", action="store_true", dest="normals",
                            help="standarize by normal expression and scale expression vals. Default only log2.  Bool")

    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # read input parameters
    pipe = pipeline("localhost", "oabalbin", "oscar", "mouseHuman")
    
    # Get the option from the input_param file
    expr_folder, seed_file, db_file, samples, bfrm, outdirs, run_param = pipe.parse_input_annotfiles(open(options.annotFile))
    
    # Create experiment folder tree and log file. 
    odir = cm.directoryTree()
    tstampedfolder, tstamp = odir.create_tstampdir(outdirs['outdir'])
    #' /bfrm, /drugsen, /drugsen/plots,/sdv' 
    dirs = odir.make_directory_tree()
    logfile = cm.logFile(tstampedfolder+'/'+tstamp+'_readme.txt') 
    
    ####### output files
    
    outfile1 = dirs[0]+tstamp+'_expression_matHsp.txt'  # Expression values
    outfile2 = dirs[0]+tstamp+'_expression_matMmus.txt'  # geneNames
    
    ######### Get raw data  ############
    # Get the database files
    #read folder for expression profiles

    
    exp_files = pipe.read_files_folder(expr_folder['expdir'],'.txt')
    
    tp = pmd.parser()
    pipe.mt.create_connector()
    gtp = gpm.geneparser()
    ######## Create database
    if options.dbtrue:
        pipe.initialize_database(exp_files,tp, options.dbtrue, options.softtrue, options.platformFile, True, True)
    #sys.exit(0)
    
    samplelist = tp.list_of_names(open(samples['samples']))
    
    if options.allgenes:
        extgenlist = gtp.list_of_genes_kgXref(open(seed_file['allgene']))  
    else:
        extgenlist = tp.list_of_names(open(seed_file['univfile']))

    pipe.mq.create_connector()
    
        ### Create the expMat array
    pvalue_th=str(1) #str(1)#
    log_fold_change=0.0
    expMat_bfrm = pipe.mq.get_exp_mat_2pval(samplelist, extgenlist, pipe.tableName, pvalue_th)
     
    # Create sample index
    expMat_bfrm.create_sample_index()
    # Normalize the expression values data. log2 and oncomine scaling, or only log2 transform.
    # or it normalize using quantile normalization
    
    '''
    if options.norm:
        expMat_bfrm.normalization()
        #expMat_bfrm.quantile_norm()
        #myplt.density_plot(dirs[3]+'/'+tstamp+'_qnorm_plot.png',expMat_bfrm.expVal.filled(fill_value=0))
    else:
        expMat_bfrm.log2_norm()
    '''
    
    # extract the genes with log fold change greater than threshold
    if abs(log_fold_change) >=0.0 :
        index = np.array(range(expMat_bfrm.expVal.shape[0]))
        mean_exp_val = np.mean(expMat_bfrm.expVal, axis=1)
        genes_to_keep = abs(mean_exp_val) >= log_fold_change
        print expMat_bfrm.expVal.shape
        expMat_bfrm.expVal = expMat_bfrm.expVal[genes_to_keep,:]
        print expMat_bfrm.expVal.shape
        
        samples = [sp[1] for sp in  expMat_bfrm.sample]
        print samples
    
        expMat_bfrm.geneSym = deque(np.array(expMat_bfrm.geneSym)[genes_to_keep])
        
        
        if options.norm:
            expMat_bfrm.normalization()
            #expMat_bfrm.quantile_norm()
            #myplt.density_plot(dirs[3]+'/'+tstamp+'_qnorm_plot.png',expMat_bfrm.expVal.filled(fill_value=0))
        else:
            expMat_bfrm.log2_norm()

        
        outoutfile1 = open(outfile1,'w')
        outoutfile1.write('geneName'+'\t'+",".join(samples).replace(',','\t')+'\n')
        write_array2file(outoutfile1, expMat_bfrm.expVal, expMat_bfrm.geneSym)
    
    else:
        
        if options.norm:
            expMat_bfrm.normalization()
            #expMat_bfrm.quantile_norm()
            #myplt.density_plot(dirs[3]+'/'+tstamp+'_qnorm_plot.png',expMat_bfrm.expVal.filled(fill_value=0))
        else:
            expMat_bfrm.log2_norm()

        
        samples = [sp[1] for sp in  expMat_bfrm.sample]
        print samples
        
        outoutfile1 = open(outfile1,'w')
        outoutfile1.write('geneName'+'\t'+",".join(samples).replace(',','\t')+'\n')
        write_array2file(outoutfile1, expMat_bfrm.expVal, expMat_bfrm.geneSym)
    
    sys.exit(0)



# Create an annotation table for the mouse -> human genes




# read the microarray file and make a small table database with it. 

# translate the mouse genes to human genes, incorporate them into the expMat array

# Create the custom matArrays

# print the table of human and mouse expression matrix after redundant extraction. 